/**
 * Copyright (C) 2022 Carnegie Mellon University
 *
 * This file is part of the TCP in the Wild course project developed for the
 * Computer Networks course (15-441/641) taught at Carnegie Mellon University.
 *
 * No part of the project may be copied and/or distributed without the express
 * permission of the 15-441/641 course staff.
 *
 *
 * This file implements the CMU-TCP backend. The backend runs in a different
 * thread and handles all the socket operations separately from the application.
 *
 * This is where most of your code should go. Feel free to modify any function
 * in this file.
 */

#include "backend.h"

#include <poll.h>
#include <stdint.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <sys/socket.h>
#include <sys/types.h>
#include <sys/time.h>
#include <time.h>
#include <unistd.h>

#include "cmu_packet.h"
#include "cmu_tcp.h"

#define MIN(X, Y) (((X) < (Y)) ? (X) : (Y))
#define ABS(X) (((X) < 0) ? (-(X)) : (X))

int timeout_ms = 100;

enum handshake{
	UNSHAKED,
	SHAKING,
	SHAKED
};

long long get_timestamp()//获取时间戳函数
{
    long long tmp;
    struct timeval tv;

    gettimeofday(&tv, NULL);
    tmp = tv.tv_sec;
    tmp = tmp * 1000;
    tmp = tmp + (tv.tv_usec / 1000);

    return tmp;
}

enum handshake handshake_stat = UNSHAKED;

void check_for_data(cmu_socket_t *sock, cmu_read_mode_t flags);

/**
 * Tells if a given sequence number has been acknowledged by the socket.
 *
 * @param sock The socket to check for acknowledgements.
 * @param seq Sequence number to check.
 *
 * @return 1 if the sequence number has been acknowledged, 0 otherwise.
 */
int has_been_acked(cmu_socket_t *sock, uint32_t seq) {
  int result;
  while (pthread_mutex_lock(&(sock->window.ack_lock)) != 0) {
  }
  result = after(sock->window.last_ack_received, seq);
  pthread_mutex_unlock(&(sock->window.ack_lock));
  return result;
}

/**
 * Updates the socket information to represent the newly received packet.
 *
 * In the current stop-and-wait implementation, this function also sends an
 * acknowledgement for the packet.
 *
 * @param sock The socket used for handling packets received.
 * @param pkt The packet data received by the socket.
 */
void handle_message(cmu_socket_t *sock, uint8_t *pkt) {
  cmu_tcp_header_t *hdr = (cmu_tcp_header_t *)pkt;
  uint8_t flags = get_flags(hdr);
  // change some code
  switch (flags) {
    case ACK_FLAG_MASK | SYN_FLAG_MASK: {
      uint32_t recv_ack = get_ack(hdr);
      if (handshake_stat==SHAKING && recv_ack==(sock->window.last_ack_received+1)) {
        printf("hand shake 2\n");
        socklen_t conn_len = sizeof(sock->conn);
        uint32_t seq = recv_ack;

        // No payload.
        uint8_t *payload = NULL;
        uint16_t payload_len = 0;

        // No extension.
        uint16_t ext_len = 0;
        uint8_t *ext_data = NULL;

        uint16_t src = sock->my_port;
        uint16_t dst = ntohs(sock->conn.sin_port);
        uint32_t ack = get_seq(hdr) + 1;
        uint16_t hlen = sizeof(cmu_tcp_header_t);
        uint16_t plen = hlen;
        uint8_t flags = ACK_FLAG_MASK;
        uint16_t adv_window = 1;
        uint8_t *response_packet = create_packet(src, dst, seq, ack, hlen, plen, flags, adv_window,
                          ext_len, ext_data, payload, payload_len);

        sock->window.last_ack_received = recv_ack;
        sock->window.next_seq_expected = ack;
        handshake_stat = SHAKED; 
        sendto(sock->socket, response_packet, plen, 0,
              (struct sockaddr *)&(sock->conn), conn_len);
        free(response_packet);

      }
      break;
    }
    case ACK_FLAG_MASK: {
      uint32_t ack = get_ack(hdr);
      uint32_t seq = get_seq(hdr);
      if (after(ack, sock->window.last_ack_received) > 0) {
        if(handshake_stat==SHAKING){
          handshake_stat = SHAKED;
        }
        sock->window.last_ack_received = ack;
        sock->window.next_seq_expected = seq;
      }
      break;
    }
    case SYN_FLAG_MASK:{
      if(handshake_stat!=UNSHAKED){
        break;
      }
      handshake_stat = SHAKING;
      printf("hand shake 1\n");

      socklen_t conn_len = sizeof(sock->conn);
      uint32_t seq = sock->window.last_ack_received;

      // No payload.
      uint8_t *payload = NULL;
      uint16_t payload_len = 0;

      // No extension.
      uint16_t ext_len = 0;
      uint8_t *ext_data = NULL;

      uint16_t src = sock->my_port;
      uint16_t dst = ntohs(sock->conn.sin_port);
      uint32_t ack = get_seq(hdr) + 1;
      sock->window.next_seq_expected = ack;
      uint16_t hlen = sizeof(cmu_tcp_header_t);
      uint16_t plen = hlen;
      uint8_t flags = ACK_FLAG_MASK | SYN_FLAG_MASK;
      uint16_t adv_window = 1;
      uint8_t *response_packet = create_packet(src, dst, seq, ack, hlen, plen, flags, adv_window,
                        ext_len, ext_data, payload, payload_len);

      sendto(sock->socket, response_packet, plen, 0,
              (struct sockaddr *)&(sock->conn), conn_len);
      free(response_packet);

      break;
    }
    default: {
      socklen_t conn_len = sizeof(sock->conn);
      uint32_t seq = sock->window.last_ack_received;
      
      // No payload.
      uint8_t *payload = NULL;
      uint16_t payload_len = 0;

      // No extension.
      uint16_t ext_len = 0;
      uint8_t *ext_data = NULL;

      uint16_t src = sock->my_port;
      uint16_t dst = ntohs(sock->conn.sin_port);
      uint32_t ack = get_seq(hdr) + get_payload_len(pkt);
      printf("seq = %d and expected = %d \n",get_seq(hdr),sock->window.next_seq_expected);
      if(get_seq(hdr) != sock->window.next_seq_expected){
        ack = sock->window.next_seq_expected;
      }
      uint16_t hlen = sizeof(cmu_tcp_header_t);
      uint16_t plen = hlen + payload_len;
      uint8_t flags = ACK_FLAG_MASK;
      uint16_t adv_window = 1;
      uint8_t *response_packet =
          create_packet(src, dst, seq, ack, hlen, plen, flags, adv_window,
                        ext_len, ext_data, payload, payload_len);

      sendto(sock->socket, response_packet, plen, 0,
             (struct sockaddr *)&(sock->conn), conn_len);
      free(response_packet);

      seq = get_seq(hdr);
      if (seq == sock->window.next_seq_expected) { 
        sock->window.next_seq_expected = seq + get_payload_len(pkt);
        payload_len = get_payload_len(pkt);
        payload = get_payload(pkt);

        // Make sure there is enough space in the buffer to store the payload.
        sock->received_buf =
            realloc(sock->received_buf, sock->received_len + payload_len);
        memcpy(sock->received_buf + sock->received_len, payload, payload_len);
        sock->received_len += payload_len;
      }
    }
  }
}

/**
 * Checks if the socket received any data.
 *
 * It first peeks at the header to figure out the length of the packet and then
 * reads the entire packet.
 *
 * @param sock The socket used for receiving data on the connection.
 * @param flags Flags that determine how the socket should wait for data. Check
 *             `cmu_read_mode_t` for more information.
 */
void check_for_data(cmu_socket_t *sock, cmu_read_mode_t flags) {
  cmu_tcp_header_t hdr;
  uint8_t *pkt;
  socklen_t conn_len = sizeof(sock->conn);
  ssize_t len = 0;
  uint32_t plen = 0, buf_size = 0, n = 0;

  while (pthread_mutex_lock(&(sock->recv_lock)) != 0) {
  }
  switch (flags) {
    case NO_FLAG:
      len = recvfrom(sock->socket, &hdr, sizeof(cmu_tcp_header_t), MSG_PEEK,
                     (struct sockaddr *)&(sock->conn), &conn_len);
      break;
    case TIMEOUT: {
      // Using `poll` here so that we can specify a timeout.
      struct pollfd ack_fd;
      ack_fd.fd = sock->socket;
      ack_fd.events = POLLIN;
      // Timeout after 3 seconds.
      if (poll(&ack_fd, 1, timeout_ms) <= 0) {
        break;
      }
    }
    // Fallthrough.
    case NO_WAIT:
      len = recvfrom(sock->socket, &hdr, sizeof(cmu_tcp_header_t),
                     MSG_DONTWAIT | MSG_PEEK, (struct sockaddr *)&(sock->conn),
                     &conn_len);
      break;
    default:
      perror("ERROR unknown flag");
  }
  if (len >= (ssize_t)sizeof(cmu_tcp_header_t)) {
    plen = get_plen(&hdr);
    pkt = malloc(plen);
    while (buf_size < plen) {
      n = recvfrom(sock->socket, pkt + buf_size, plen - buf_size, 0,
                   (struct sockaddr *)&(sock->conn), &conn_len);
      buf_size = buf_size + n;
    }
    handle_message(sock, pkt);
    free(pkt);
  }
  pthread_mutex_unlock(&(sock->recv_lock));
}

void init_send(void* in){
  if(handshake_stat!=UNSHAKED){
    return;
  }
  handshake_stat = SHAKING;
  printf("hand shake start\n");
  cmu_socket_t *sock = (cmu_socket_t *)in;
  uint8_t *msg;
  size_t conn_len = sizeof(sock->conn);

  int sockfd = sock->socket;

  uint16_t payload_len = 0;
  uint16_t src = sock->my_port;
  uint16_t dst = ntohs(sock->conn.sin_port);
  uint32_t seq = sock->window.last_ack_received;
  uint32_t ack = sock->window.next_seq_expected;
  uint16_t hlen = sizeof(cmu_tcp_header_t);
  uint16_t plen = hlen;
  uint8_t flags = SYN_FLAG_MASK;
  uint16_t adv_window = 1;
  uint16_t ext_len = 0;
  uint8_t *ext_data = NULL;
  uint8_t *payload = NULL;

  msg = create_packet(src, dst, seq, ack, hlen, plen, flags, adv_window,
                      ext_len, ext_data, payload, payload_len);

  sendto(sockfd, msg, plen, 0, (struct sockaddr *)&(sock->conn), conn_len);
}

/**
 * Breaks up the data into packets and sends a single packet at a time.
 *
 * You should most certainly update this function in your implementation.
 *
 * @param sock The socket to use for sending data.
 * @param data The data to be sent.
 * @param buf_len The length of the data being sent.
 */
void single_send(cmu_socket_t *sock, uint8_t *data, int buf_len) {
  uint8_t *msg;
  uint8_t *data_offset = data;
  size_t conn_len = sizeof(sock->conn);

  int sockfd = sock->socket;
  if (buf_len > 0) {
    while (buf_len != 0) {
      uint16_t payload_len = MIN((uint16_t)buf_len, (uint16_t)MSS);

      uint16_t src = sock->my_port;
      uint16_t dst = ntohs(sock->conn.sin_port);
      uint32_t seq = sock->window.last_ack_received;
      uint32_t ack = sock->window.next_seq_expected;
      uint16_t hlen = sizeof(cmu_tcp_header_t);
      uint16_t plen = hlen + payload_len;
      uint8_t flags = 0;
      uint16_t adv_window = 1;
      uint16_t ext_len = 0;
      uint8_t *ext_data = NULL;
      uint8_t *payload = data_offset;

      msg = create_packet(src, dst, seq, ack, hlen, plen, flags, adv_window,
                          ext_len, ext_data, payload, payload_len);
      buf_len -= payload_len;

      while (1) {
        // FIXME: This is using stop and wait, can we do better?
        sendto(sockfd, msg, plen, 0, (struct sockaddr *)&(sock->conn),
               conn_len);
        check_for_data(sock, TIMEOUT);
        if (has_been_acked(sock, seq)) {
          break;
        }
      }
      data_offset += payload_len;
    }
  }
}
// add the function to send multiple packets at a time
/**
 * Breaks up the data into packets and sends multiple packages at a time.
 *
 * @param sock The socket to use for sending data.
 * @param data The data to be sent.
 * @param buf_len The length of the data being sent.
 */
void multi_send(cmu_socket_t *sock, uint8_t *data, int buf_len) {
  int window_size = WINDOW_INITIAL_WINDOW_SIZE / MSS;
  uint8_t *msg[window_size];
  uint32_t seqlist[window_size];
  uint16_t pkt_len_list[window_size];
  int retrans_flag = 0, last_flag = 0;
  long long timestamps[window_size];

  uint8_t *data_offset = data;
  size_t conn_len = sizeof(sock->conn);
  int sockfd = sock->socket;

  int esti_rtt = 0.1, dev_rtt = 0.025;

  if (buf_len > 0) {
    int base = 0;
    int next_seq_num = 0;
    uint32_t last_ack_received = sock->window.last_ack_received;
    uint32_t next_seq_expected = sock->window.next_seq_expected;
    while (buf_len != 0 || base < next_seq_num) {
      
      if (retrans_flag)
      { 
        for(int i = base;i < next_seq_num;i++){
          sendto(sockfd, msg[i % window_size], pkt_len_list[i % window_size], 0, (struct sockaddr *)&(sock->conn),
                conn_len);
          check_for_data(sock, TIMEOUT);
        }
        base++;
        retrans_flag = 0;
        continue;
      }
      while (buf_len >= MSS && (next_seq_num - base) < window_size)
      {
        uint16_t src = sock->my_port;
        // printf("base is %d and next is %d\n",base,next_seq_num);
        uint16_t dst = ntohs(sock->conn.sin_port);
        uint32_t seq = last_ack_received;
        printf("seq is %d, buflen: %d\n",seq, buf_len);
        // printf("seq is %d\n",seq);
        uint32_t ack = next_seq_expected;
        last_ack_received += MSS;
        uint16_t hlen = sizeof(cmu_tcp_header_t);
        uint16_t plen = hlen + MSS;
        uint8_t flags = 0;
        uint16_t adv_window = 1;
        uint16_t ext_len = 0;
        uint8_t *ext_data = NULL;
        uint8_t *payload = data_offset;
        msg[next_seq_num % window_size] = create_packet(src, dst, seq, ack, hlen, plen, flags, adv_window,
                                      ext_len, ext_data, payload, MSS);
        seqlist[next_seq_num % window_size] = seq;
        timestamps[next_seq_num % window_size] = get_timestamp();
        pkt_len_list[next_seq_num % window_size] = plen;
        sendto(sockfd, msg[next_seq_num % window_size], plen, 0, (struct sockaddr *)&(sock->conn),
                conn_len);
        next_seq_num++;
        buf_len -= MSS;
        data_offset += MSS;
      }
      if ((next_seq_num - base < window_size) && buf_len > 0)
      {
        uint16_t src = sock->my_port;
        uint16_t dst = ntohs(sock->conn.sin_port);
        uint32_t seq = last_ack_received;
        printf("last seq is %d, buflen: %d\n",seq, buf_len);
        last_flag = 1;
        last_ack_received += buf_len;
        uint32_t ack = next_seq_expected;
        uint16_t hlen = sizeof(cmu_tcp_header_t);
        uint16_t plen = hlen + buf_len;
        uint8_t flags = 0;
        uint16_t adv_window = 1;
        uint16_t ext_len = 0;
        uint8_t *ext_data = NULL;
        uint8_t *payload = data_offset;
        msg[next_seq_num % window_size] = create_packet(src, dst, seq, ack, hlen, plen, flags, adv_window,
                                          ext_len, ext_data, payload, buf_len);
        seqlist[next_seq_num % window_size] = seq;
        pkt_len_list[next_seq_num % window_size] = plen;                                 
        sendto(sockfd, msg[next_seq_num % window_size], plen, 0, (struct sockaddr *)&(sock->conn),
              conn_len);
        next_seq_num++;
        data_offset += buf_len;
        buf_len = 0;
      }

      check_for_data(sock, TIMEOUT);
      int ack_num = has_been_acked(sock, seqlist[base % window_size]);
      printf("ack_num: %d, last: %d, seq: %d\n", ack_num, sock->window.last_ack_received, seqlist[base % window_size]);
      
      if (ack_num > 0) {
        if(last_flag==0){
          int real_rtt = (int)(get_timestamp() - timestamps[base % window_size]);
          dev_rtt = (int)((float)ABS(real_rtt-esti_rtt) * 0.25 + (float)dev_rtt * 0.75);
          esti_rtt = (int)((float)real_rtt * 0.125 + (float)esti_rtt * 0.875);
          timeout_ms = esti_rtt + 4*dev_rtt;
          printf("real_rtt: %d, esti_rtt: %d, dev_rtt: %d, interval: %d\n", real_rtt, esti_rtt, dev_rtt, timeout_ms);
        }
        if ((ack_num % MSS) > 0)
        {
          base = base + (ack_num / MSS) + 1;
        }
        else base += ack_num / MSS;
      }
      else{
        if(last_flag==0){
          esti_rtt *= 2;
          timeout_ms = esti_rtt + 4*dev_rtt;
          printf("timeout! esti_rtt: %d, interval: %d\n", esti_rtt, timeout_ms);
        }
        if (ack_num == 0)
        {
          retrans_flag = 1;
        }
      } 
    }
  }
}

void *begin_backend(void *in) {
  cmu_socket_t *sock = (cmu_socket_t *)in;
  int death, buf_len, send_signal;
  uint8_t *data;

  if(sock->type == TCP_INITIATOR){
    init_send(in);
  }
  while(handshake_stat != SHAKED){
    check_for_data(sock, NO_WAIT);
  }
  while (1) {
    while (pthread_mutex_lock(&(sock->death_lock)) != 0) {
    }
    death = sock->dying;
    pthread_mutex_unlock(&(sock->death_lock));

    while (pthread_mutex_lock(&(sock->send_lock)) != 0) {
    }
    buf_len = sock->sending_len;
    if (death && (buf_len == 0)) {
      break;
    }
    if (buf_len > 0) {
      data = malloc(buf_len);
      memcpy(data, sock->sending_buf, buf_len);
      sock->sending_len = 0;
      free(sock->sending_buf);
      sock->sending_buf = NULL;
      pthread_mutex_unlock(&(sock->send_lock));
      multi_send(sock, data, buf_len);
      free(data);
    } else {
      pthread_mutex_unlock(&(sock->send_lock));
    }
    check_for_data(sock, NO_WAIT);

    while (pthread_mutex_lock(&(sock->recv_lock)) != 0) {
    }
    send_signal = sock->received_len > 0;
    
    pthread_mutex_unlock(&(sock->recv_lock));

    if (send_signal) {
      pthread_cond_signal(&(sock->wait_cond));
    }
  }

  pthread_exit(NULL);
  return NULL;
}
